# 构建数据集

from paddleseg.datasets import Dataset
import paddleseg.transforms as T
from src.config.config import SIZE, DATA_ROOT, NUM_CLASSES, TRAIN_PATH, VAL_PATH, TEST_PATH

# The transforms must be a list!
# 构建训练集
train_transforms = [
    T.Resize(target_size=SIZE),
    T.RandomHorizontalFlip(),
    T.Normalize()
]

# 构建验证集
val_transforms = [
    T.Resize(target_size=SIZE),
    T.Normalize()
]

# 构建测试集
# 用predict API测试时，一定要是T.Compose类型
test_transforms_compose = T.Compose([
    T.Resize(target_size=SIZE),
    T.Normalize()
])

# 这个数据增强一定要是list类型！！！ 上面用predict API测试时是T.Compose类型
test_transforms = [
    T.Resize(target_size=SIZE),
    T.Normalize()
]

train_dataset = Dataset(
    mode='train',  # 训练模式
    dataset_root=DATA_ROOT,
    transforms=train_transforms,
    num_classes=NUM_CLASSES,
    train_path=TRAIN_PATH,
)

val_dataset = Dataset(
    mode='val',  # 验证模式
    dataset_root=DATA_ROOT,
    transforms=val_transforms,
    num_classes=NUM_CLASSES,
    val_path=VAL_PATH,
)

test_dataset = Dataset(
    mode='val',  # 评估模式
    dataset_root=DATA_ROOT,
    transforms=test_transforms,
    num_classes=NUM_CLASSES,
    val_path=TEST_PATH,  # 评估测试集
)
